# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/9/26
# @Function: 阈值密码学方案实现

import base64
import random
import hashlib
import json
from gmssl import sm2, sm4
from typing import List, Dict, Tuple, Any


class ThresholdCrypto:
    """
    阈值密码学方案实现
    允许将密钥分割成多份，只有当足够数量的密钥持有者合作时，才能解密数据
    """
    def __init__(self):
        """初始化阈值密码学工具"""
        self.sm4_key = None  # 用于加密实际数据的对称密钥
    
    def generate_sm4_key(self) -> bytes:
        """生成随机SM4密钥"""
        return bytes([random.randint(0, 255) for _ in range(16)])
    
    def split_key(self, key: bytes, n: int, t: int) -> List[Dict[str, Any]]:
        """
        使用Shamir's Secret Sharing算法将密钥分割成n份，需要至少t份才能重构
        
        Args:
            key: 要分割的密钥
            n: 分割成的份数
            t: 重构所需的最小份数
            
        Returns:
            List[Dict]: 包含密钥分片的列表
        """
        if t > n:
            raise ValueError("重构所需的份数不能大于总份数")
        
        # 将密钥转换为整数
        key_int = int.from_bytes(key, byteorder='big')
        
        # 生成多项式的系数，a_0 = 密钥，其他系数随机
        coefficients = [key_int]
        # 生成t-1个随机系数
        for _ in range(t - 1):
            coefficients.append(random.randint(0, 2**128))
        
        # 生成n个密钥分片
        shares = []
        for i in range(1, n + 1):  # x坐标从1开始
            # 计算多项式值 f(x) = a_0 + a_1*x + a_2*x^2 + ... + a_{t-1}*x^{t-1}
            y = 0
            for j, coef in enumerate(coefficients):
                y = (y + coef * (i ** j)) % (2**128)
            
            # 将分片信息保存为字典
            share = {
                'x': i,
                'y': y,
                'threshold': t,
                'total': n
            }
            shares.append(share)
        
        return shares
    
    # 修改reconstruct_key方法
    def reconstruct_key(self, shares: List[Dict[str, Any]]) -> bytes:
        """
        使用Shamir's Secret Sharing算法从密钥分片重构密钥
        
        Args:
            shares: 密钥分片列表
            
        Returns:
            bytes: 重构的密钥
        """
        if not shares:
            raise ValueError("没有提供密钥分片")
        
        # 检查是否有足够的分片
        threshold = shares[0]['threshold']
        if len(shares) < threshold:
            raise ValueError(f"需要至少{threshold}个密钥分片才能重构密钥，但只提供了{len(shares)}个")
        
        # 使用拉格朗日插值重构密钥
        secret = 0
        prime = 2**128  # 使用一个大素数作为模数
        
        # 只使用前threshold个分片
        shares = shares[:threshold]
        
        for i, share_i in enumerate(shares):
            xi = share_i['x']
            yi = share_i['y']
            
            # 计算拉格朗日基本多项式的值
            lagrange_basis = 1
            
            for j, share_j in enumerate(shares):
                if i == j:
                    continue
                
                xj = share_j['x']
                # 计算 (x - x_j) / (x_i - x_j) 在 x=0 处的值
                # 当 x=0 时，分子为 -x_j
                numerator = (-xj) % prime
                denominator = (xi - xj) % prime
                
                # 计算模逆元
                def mod_inverse(a, m):
                    g, x, y = extended_gcd(a, m)
                    if g != 1:
                        raise ValueError("模逆不存在")
                    else:
                        return x % m
                    
                def extended_gcd(a, b):
                    if a == 0:
                        return b, 0, 1
                    else:
                        gcd, x, y = extended_gcd(b % a, a)
                        return gcd, y - (b // a) * x, x
                
                inv = mod_inverse(denominator, prime)
                lagrange_basis = (lagrange_basis * numerator * inv) % prime
            
            # 将该点的贡献加到密钥中
            term = (yi * lagrange_basis) % prime
            secret = (secret + term) % prime
        
        # 将整数转换回字节
        return secret.to_bytes(16, byteorder='big')
    
    # 在ThresholdCrypto类中修改encrypt_data方法
    def encrypt_data(self, data: str, n: int, t: int) -> Dict[str, Any]:
        """
        加密数据并分割密钥
        
        Args:
            data: 要加密的数据
            n: 密钥分割的份数
            t: 重构所需的最小份数
            
        Returns:
            Dict: 包含加密数据和密钥分片的字典
        """
        # 生成SM4对称密钥
        self.sm4_key = self.generate_sm4_key()
        
        # 使用SM4加密数据
        sm4_crypt = sm4.CryptSM4()
        sm4_crypt.set_key(self.sm4_key, sm4.SM4_ENCRYPT)
        
        # 确保数据长度是16的倍数（SM4块大小），添加PKCS#7填充
        data_bytes = data.encode('utf-8')
        padding_length = 16 - (len(data_bytes) % 16)
        if padding_length < 16:
            data_bytes += bytes([padding_length]) * padding_length
        
        encrypted_data = sm4_crypt.crypt_ecb(data_bytes)
        
        # 分割密钥
        key_shares = self.split_key(self.sm4_key, n, t)
        
        # 构建返回结果
        result = {
            'encrypted_data': base64.b64encode(encrypted_data).decode('utf-8'),
            'key_shares': key_shares,
            'threshold': t,
            'total_shares': n
        }
        
        return result
    
    # 修改decrypt_data方法
    def decrypt_data(self, encrypted_package: Dict[str, Any], shares: List[Dict[str, Any]]) -> str:
        """
        使用密钥分片解密数据
        
        Args:
            encrypted_package: 包含加密数据的字典
            shares: 密钥分片列表
            
        Returns:
            str: 解密后的数据
        """
        try:
            # 重构密钥
            reconstructed_key = self.reconstruct_key(shares)
            
            # 使用重构的密钥解密数据
            encrypted_data = base64.b64decode(encrypted_package['encrypted_data'])
            sm4_crypt = sm4.CryptSM4()
            sm4_crypt.set_key(reconstructed_key, sm4.SM4_DECRYPT)
            decrypted_data = sm4_crypt.crypt_ecb(encrypted_data)
            
            # 移除PKCS#7填充
            if decrypted_data:
                padding_length = decrypted_data[-1]
                if padding_length < 16 and all(b == padding_length for b in decrypted_data[-padding_length:]):
                    decrypted_data = decrypted_data[:-padding_length]
            
            return decrypted_data.decode('utf-8')
        except Exception as e:
            print(f"解密过程中出错: {e}")
            # 打印更多调试信息
            print(f"重构的密钥: {reconstructed_key.hex() if 'reconstructed_key' in locals() else 'N/A'}")
            if 'decrypted_data' in locals() and decrypted_data:
                print(f"解密数据的前20字节: {decrypted_data[:20].hex()}")
            raise


def demo():
    """演示阈值密码学方案的使用"""
    # 创建阈值密码学工具
    crypto = ThresholdCrypto()
    
    # 数据拥有者：叶磊老师
    data = "这是叶磊老师的敏感研究数据，需要安全共享给徐博老师和初京刚老师"
    
    # 设置阈值参数：总共分成3份，需要至少2份才能解密
    n = 3  # 总份数
    t = 2  # 阈值
    
    try:
        # 加密数据并分割密钥
        encrypted_package = crypto.encrypt_data(data, n, t)
        
        print("加密数据包:")
        print(json.dumps(encrypted_package, indent=2))
        
        # 分发密钥分片给三位老师
        ye_lei_share = encrypted_package['key_shares'][0]
        xu_bo_share = encrypted_package['key_shares'][1]
        chu_jinggang_share = encrypted_package['key_shares'][2]
        
        print("\n叶磊老师的密钥分片:")
        print(json.dumps(ye_lei_share, indent=2))
        
        print("\n徐博老师的密钥分片:")
        print(json.dumps(xu_bo_share, indent=2))
        
        print("\n初京刚老师的密钥分片:")
        print(json.dumps(chu_jinggang_share, indent=2))
        
        # 场景1：叶磊老师和徐博老师合作解密
        print("\n场景1：叶磊老师和徐博老师合作解密")
        try:
            shares = [ye_lei_share, xu_bo_share]
            decrypted_data = crypto.decrypt_data(encrypted_package, shares)
            print(f"解密结果: {decrypted_data}")
            assert decrypted_data == data, "解密结果与原始数据不匹配"
        except Exception as e:
            print(f"解密失败: {e}")
        
        # 场景2：徐博老师和初京刚老师合作解密
        print("\n场景2：徐博老师和初京刚老师合作解密")
        try:
            shares = [xu_bo_share, chu_jinggang_share]
            decrypted_data = crypto.decrypt_data(encrypted_package, shares)
            print(f"解密结果: {decrypted_data}")
            assert decrypted_data == data, "解密结果与原始数据不匹配"
        except Exception as e:
            print(f"解密失败: {e}")
        
        # 场景3：只有叶磊老师一个人尝试解密（应该失败）
        print("\n场景3：只有叶磊老师一个人尝试解密（应该失败）")
        try:
            shares = [ye_lei_share]
            decrypted_data = crypto.decrypt_data(encrypted_package, shares)
            print(f"解密结果: {decrypted_data}")
        except Exception as e:
            print(f"解密失败（符合预期）: {e}")
    except Exception as e:
        print(f"程序执行出错: {e}")


if __name__ == "__main__":
    demo()